package com.yifeng.repo.toolkit.mq;

import com.google.common.base.Strings;
import org.apache.rocketmq.client.producer.MessageQueueSelector;
import org.apache.rocketmq.common.message.Message;
import org.apache.rocketmq.common.message.MessageConst;
import org.apache.rocketmq.common.message.MessageQueue;
import org.springframework.messaging.MessageHeaders;

import java.util.List;

/**
 * Created by daibing on 2022/9/25.
 */
public class OrderlyMessageQueueSelector implements MessageQueueSelector {
    private static final String ROUTING_KEY = "x-routing-key";

    @Override
    public MessageQueue select(List<MessageQueue> mqs, Message msg, Object arg) {
        // 优先使用routingKey
        String routingKey = (String) ((MessageHeaders) arg).get(ROUTING_KEY);
        if (!Strings.isNullOrEmpty(routingKey)) {
            int index = routingKey.hashCode() % mqs.size();
            return mqs.get(index);
        }

        // 如果没有设置routingKey，就使用tag
        String tag = (String) ((MessageHeaders) arg).get(MessageConst.PROPERTY_TAGS);
        if (Strings.isNullOrEmpty(tag)) {
            throw new IllegalArgumentException("Not found tag from headers");
        }
        // 同一个tag的消息进入同一个队列，就可以保证tag内顺序消费
        int index = (tag.hashCode() & Integer.MAX_VALUE) % mqs.size();
        return mqs.get(index);
    }
}
